import os
from langchain_deepseek import ChatDeepSeek
from langchain_core.prompts import ChatPromptTemplate
# from prompt import MEDICAL_PROMPT_TEMPLATE

# 初始化DeepSeek模型
def init_llm():
    return ChatDeepSeek(
        api_key=os.getenv("DEEPSEEK_API_KEY"),
        model="deepseek-chat",
        temperature=0.1,
        max_tokens=1000,
        max_retries=3,
        timeout=30,
        verbose=True
    )

# 生成就医建议的提示模板
MEDICAL_PROMPT_TEMPLATE = """
您是一位神经内科主任医师，请根据以下脑部影像预测结果生成专业就医建议：
- 预测类别: {label} ({label_cn})
- 置信度: {confidence:.2%}
- 概率分布: 类别0={prob_0:.2%}, 类别1={prob_1:.2%}

请包含以下内容：
1. 结果解读（非技术术语）
2. 建议的进一步检查
3. 推荐就诊科室
4. 日常护理建议

重要：避免诊断结论，只提供建议。结尾注明"具体诊疗请遵医嘱"。
"""

def generate_medical_advice(prediction: dict):
    llm = init_llm()
    label_cn = "认知正常(CN)" if prediction['label'] == 'CN' else "阿尔茨海默病(AD)"
    
    prompt = ChatPromptTemplate.from_template(MEDICAL_PROMPT_TEMPLATE)
    chain = prompt | llm
    
    return chain.invoke({
        "label": prediction['label'],
        "label_cn": label_cn,
        "confidence": prediction['confidence'],
        "prob_0": prediction['probabilities']['0'],
        "prob_1": prediction['probabilities']['1']
    }).content